import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread


class PEC(HasTraits):
    nz = Int(100)
    s = Float(.5)
    nd = Int(20)
    nt = Int(5)
    stepNumber = Int(400)
    parameter_view = View(
        Item('nz', label='网格数量'),
        Item('s', label='时间步长因子'),
        Item('nd', label='空间步长因子'),
        Item('nt', label='曲线更新步数'),
        Item('stepNumber', label='时间步数'),
    )

class PML(HasTraits):
    nz = Int(100)
    s = Float(.5)
    nd = Int(20)
    na = Int(10)
    nt = Int(5)
    stepNumber = Int(400)
    parameter_view = View(
        Item('nz', label='网格数量'),
        Item('s', label='时间步长因子'),
        Item('nd', label='空间步长因子'),
        Item('nt', label='曲线更新步数'),
        Item('stepNumber', label='时间步数'),
        Item('na', label='吸收层厚度'),
    )

class Form(HasTraits):
    np.seterr(divide='ignore')

    paraPEC = Instance(PEC, ())
    paraPML = Instance(PML, ())
    method = Enum('pec', 'pml')
    plot = Instance(Plot)
    progress = Str('0%')
    thread = QThread()

    def run(self):
        match self.method:
            case 'pec':
                self.pec()
            case 'pml':
                self.pml()
    def pec(self):
        nz = self.paraPEC.nz
        s = self.paraPEC.s
        nd = self.paraPEC.nd
        nt = self.paraPEC.nt
        e = np.zeros(nz + 1)
        h = 0
        for n in range(self.paraPEC.stepNumber + 1):
            h -= s * np.diff(e)
            e[1:-1] -= s * np.diff(h)
            e[0] = np.sin(2 * np.pi * n * s / nd)
            if n % nt == 0:
                self.plot.data.update_data(
                    x=np.arange(nz),
                    e=e,
                )
                self.plot.title = f'Step number: {n}/{self.paraPEC.stepNumber}'
                time.sleep(.1)
                self.progress = f'{n / self.paraPEC.stepNumber:<.1%}'

    def pml(self):
        nz = self.paraPML.nz
        s = self.paraPML.s
        nd = self.paraPML.nd
        nt = self.paraPML.nt
        na = self.paraPML.na
        e = np.zeros(nz + 1)
        h = 0
        ze = 0
        zh = 0
        ae = np.exp(-4 * s * (np.arange(1, na) / na) ** 4)
        am = np.exp(-4 * s * (np.arange(.5, na) / na) ** 4)
        be = 1 - ae
        bm = 1 - am
        for n in range(self.paraPML.stepNumber + 1):
            h -= s * np.diff(e)
            h[-na:] -= s * ze
            zh = ae * zh - be * np.diff(h[-na:])
            e[1:-1] -= s * np.diff(h)
            e[-na:-1] -= s * zh
            ze = am * ze - bm * np.diff(e[-na - 1:])
            e[0] = np.sin(2 * np.pi * n * s / nd)
            if n % nt == 0:
                self.plot.data.update_data(
                    x=np.arange(nz),
                    e=e,
                )
                self.plot.title = f'Step number: {n}/{self.paraPML.stepNumber}'
                time.sleep(.1)
                self.progress = f'{n / self.paraPML.stepNumber:<.1%}'

    def traits_init(self):
        self.thread.run = self.run
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            x=[],
            e=[],
        )
        self.plot.x_axis.title = "x position"
        self.plot.y_axis.title = "Electric / (V/m)"
        self.plot.plot(('x', 'e'), name='E', color='auto')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)
        self.plot.tools.append(
            PanTool(self.plot, constrain_key="shift")
        )
        self.plot.overlays.append(
            ZoomTool(component=self.plot, tool_mode="box", always_on=False)
        )

    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        HGroup(
            Item('progress', label='计算进度'),
            spring,
        ),
        width=1200,
        height=900,
        resizable=True,
    )
    parameter_view = View(
        HGroup(
            Item('method', style='custom'),
            spring
        ),
        UItem('paraPEC', style='custom', visible_when='method=="pec"'),
        UItem('paraPML', style='custom', visible_when='method=="pml"'),
        resizable=True,
        buttons=OKCancelButtons,
    )

if __name__ == "__main__":
    form = Form()
    form.configure_traits(view='parameter_view')
    form.thread.start()
    form.configure_traits()
